// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================


/*!
\file
\brief	Main Body of the PMODEL MPI Program
*/
/*! 
\defgroup	pmodel PMODEL - Allelic Model Association Tests for Binary Trait
\ingroup    mpi_stats
\brief		PMODEL - Allelic %Model Association Tests for Binary Trait
\details
<p>
This program performs model association tests for SNPs and a binary trait.<br>
A model is specified by selecting the appropriate command-line %option.</p>

<p>PMODEL runs in 2 modes</p>
<dl>
<dt>Mode 1</dt>
<dd>Multi-core MPI invocations which perform statistical calculation with permutation.</dd>
<dt>Mode 2</dt>
<dd>Single-core MPI call which parses the permutation data and creates the final output file.</dd>
</dl>

<p>By default PMODEL runs in <em>Mode 1</em>, however if the <em>--merge</em> parameter is specified, the 
program runs in the <em>Mode 2</em> where it attempts to read the permutation data.<br>
If multiple cores are assigned to a run when <em>--merge</em> is specified, PMODEL exits with an error code.
</p>
<p>The statistical output is placed in a file with the suffix \ref SUFFIX_MODEL_FILE .</p>
<p>If successful PMODEL returns with an exitcode of 0, if it encountered an error the programs returns 
with an exit code of -1.</p>

<h2>OPTIONS</H2>

<table class=options border=0>
<tr><th class=param_header colspan=1>PARAMETER</th>
<th class=param_header colspan=3>DESCRIPTION</th></tr>

<tr><td class=param_name colspan=1>--id STRING</td>
<td class=param_desc colspan=3>Unique identifer for a cluster job (also sets output file stem).</td></tr>

<tr><td class=param_name colspan=1>--rdir PATH</td>
<td class=param_desc colspan=3>Path to remote working directory on central file system.</td></tr>

<tr><td class=param_name colspan=1>--ldir PATH</td>
<td class=param_desc colspan=3>Path to local working directory on each core of the cluster.</td></tr>

<tr><td class=param_name colspan=1>--bmap FILE</td>
<td class=param_desc colspan=3>Path to the BMAP file (binary-encoded SNP list).</td></tr>

<tr><td class=param_name colspan=1>--samples FILE</td>
<td class=param_desc colspan=3>Path to the Sample file (people and phenotype).</td></tr>

<tr><td class=param_name colspan=1>--sped FILE</td>
<td class=param_desc colspan=3>Path to the SPED file (Allelic data, SNP Major, bit-strings).</td></tr>

<tr><td class=param_name colspan=1>--perm FILE</td>
<td class=param_desc colspan=3>Path to the permutation indices file.</td></tr>

<tr><td class=param_name colspan=1>--merge</td>
<td class=param_desc colspan=3>Signal PMODEL to merge permutation results and create final output file (single-core operation).</td></tr>

<tr><td class=param_name colspan=1>--missed-rank 0..101</td>
<td class=param_desc colspan=3>%Set processor rank to number specified by this parameter instead of rank assigned by MPI.<br>
To be used only if a section of a calculation needs to run a single core if a processor failed.
</td></tr>

<tr><td class=param_name colspan=1>--missed-procs 1..102</td>
<td class=param_desc colspan=3>%Set the No. processors associated with a calculation incase processor rank set by <em>missed-rank</em>.</td></tr>

<tr><td class=param_name colspan=1>--dominant</td>
<td class=param_desc colspan=3>Run dominant gene action (1df) test.</td></tr>

<tr><td class=param_name colspan=1>--recessive</td>
<td class=param_desc colspan=3>Run recessive gene action (1df) test.</td></tr>

<tr><td class=param_name colspan=1>--trend</td>
<td class=param_desc colspan=3>Run Cochran-Armitage trend test.</td></tr>

<tr><td class=param_name colspan=1>--genotypic</td>
<td class=param_desc colspan=3>Run Genotypic (2 df) test.</td></tr>

<tr><td class=param_name colspan=1>--allelic</td>
<td class=param_desc colspan=3>Run basic allelic test.</td></tr>

<tr><td class=param_name colspan=1>--ci 0..1.0</td>
<td class=param_desc colspan=3>%Set the confidence interval (CI) of the Odds-Ratio (OR) for the dominant, recessive, trend and allelic tests.<br>
Default ci=0.95
</td></tr>

<tr><td class=param_name colspan=1>--sort-by COLUMN_NAME</td>
<td class=param_desc colspan=3>Sort column for the final CSV file containing the statistical output.<br>
Default sort column=EMP1.
</td></tr>
</table>
*/

#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include "SMG_All.h"

/*!
\brief Command-line options of PMODEL
*/
struct option long_options[] = {
	{ "bmap", required_argument, NULL, 'm' }, 
	{ "samples", required_argument, NULL, 's' }, // Path to the map file.
	{ "sped", required_argument, NULL, 'p' }, // Path to the PED file.
	{ "id", required_argument, NULL, 'i' }, // Run ID
	{ "rdir", required_argument, NULL, 'r'}, // Remote Output direct
	{ "ldir", required_argument, NULL, 'l'}, // Local Output direct
	{ "perm", required_argument, NULL, 'P'}, // Path to the permutation file.
	{ "merge", optional_argument, NULL, 'M'}, // Flag about to do the merge.
	{ "missed-rank", required_argument, NULL, 'n'}, // A rank that failed for whatever reason that needs re-running.
	{ "missed-procs", required_argument, NULL, 'o'}, // No. processors associated with the original calculation.
	{ "dominant", optional_argument, NULL, 'D' }, // Calculate dominant
	{ "recessive", optional_argument, NULL, 'R' }, // Calculate recessive
	{ "trend", optional_argument, NULL, 'T' }, // Calculate trend
	{ "genotypic", optional_argument, NULL, 'G' }, // Calculate genotypic
	{ "allelic", optional_argument, NULL, 'A' }, // Calculate Allelic
	{ "ci", required_argument, NULL, 'c' }, // Confidence interval
	{ "sort-by", required_argument, NULL, 'S' }, // Column Name to sort final output by.
	{ 0,    0,    0,    0   } 
};

/*!
\brief Main body of the PMODEL program
\ingroup pmodel
*/
int pmodel_main(int argc, char *argv[]) {
	char ch, *spedData = NULL;
	struct selected_options ops;
	int myrank, numprocs, result;
	BOOL jobDone = FALSE;
	struct chromosomes chromos;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	int nArgCount = 0, nRequiredArgs = 7, i, j, nAdj = 0;
	struct small_locus *map = NULL;
	struct sample *samples = NULL;
	int nMap = 0, nSamples = 0, nSpedSamples = 0, *pPermutations = NULL;
	int nSNPs = 0, nSpedRecordSizeInBytes = 0;
	BOOL isDoingMerge = FALSE;
	int perm_rows = 0, perm_columns = 0;
	double *originals = NULL, *results = NULL, zt;
	int *R = NULL, *maxR = NULL, missedRank = -1, missedProcs = -1, nModels = 0;
	struct adjusted *adj = NULL;

	// Initialise MPI environment, get size and rank.
	if ((result = MPI_Init(&argc,&argv)) != MPI_SUCCESS) 
		mpi_error(result, "Initialisation");
	if ((result = MPI_Comm_size(MPI_COMM_WORLD,&numprocs)) != MPI_SUCCESS)
		mpi_error(result, "Number of processors");
	if ((result = MPI_Comm_rank(MPI_COMM_WORLD,&myrank)) != MPI_SUCCESS)
		mpi_error(result, "Processor number");

	init_genrand((unsigned long) time(NULL));
	options_init(&ops);
	chromosomes_init(&chromos);
	chromosomes_define_human(&chromos);

	// Want to set these set via the command-line.
	*ops.ldir = *ops.rdir = *ops.szRunId = '\0'; 
	ops.verbose = FALSE;

	while ((ch = (char) getopt_long(argc, argv, "m:v:p:f:i:", long_options, NULL)) != -1) {
		double value;
		switch (ch) {
			case 'i':
				strcpy(ops.szRunId, optarg);
				nArgCount++;
				break;
			case 's': // Sample File
				strcpy(ops.szSampleFileName, optarg);
				nArgCount++;
				break;
			case 'p': // SPED file
				strcpy(ops.szSpedFileName, optarg);
				nArgCount++;
				break;
			case 'm': // Bmap File
				strcpy(ops.szBmapFileName, optarg);
				nArgCount++;
				break;
			case 'r': // Remote output directory
				strcpy(ops.rdir, optarg);
				nArgCount++;
				break;
			case 'l': // Local output directory
				strcpy(ops.ldir, optarg);
				nArgCount++;
				break;
			case 'P': // Path to permutation file.
				strcpy(ops.szPermFileName, optarg);
				nArgCount++;
				break;
			case 'M':
				isDoingMerge = TRUE;
				break;
			case 'n': // Missed rank to re-run calculation on a single processor.
				missedRank = atoi(optarg);
				break;
			case 'o': // No. processors associated with a 'missed' rank.
				missedProcs = atoi(optarg);
				break;
			case 'R': // Do recessive calculation.
				ops.model_perm_rec = TRUE;
				nModels++;
				break;
			case 'T': // Do trend calculation.
				ops.model_perm_trend = TRUE;
				nModels++;
				break;
			case 'G': // Do genotypic calculation.
				ops.model_perm_gen = TRUE;
				nModels++;
				break;
			case 'A': // Do genotypic calculation.
				ops.model_perm_allelic = TRUE;
				nModels++;
				break;
			case 'D': // Do genotypic calculation.
				ops.model_perm_dom = TRUE;
				nModels++;
				break;
			case 'S': // Sort column name
				if (strlen(optarg) > 0) strcpy(ops.szSortColumnName, optarg);
				break;
			case 'c': // The confidence interval
				value = atof(optarg);
				if (value >= 0.5 && value <= 1.0) {
					ops.ci_level = value;
				}
				break;
			default:
				break;
		}
	}
	zt = ltqnorm(1 - (1 - ops.ci_level) / 2); // For the CI of the OR

	// If the required arguements are not specified, there's no point carrying on.
	if (isDoingMerge) nRequiredArgs = 3; // just need ID, LDIR & RDIR
	if (nArgCount != nRequiredArgs || (nModels == 0 && !isDoingMerge) || nModels > 1) {
		chromosomes_free_attributes(&chromos);
		MPI_Finalize();
		return -1;
	}

	// See if a calculation needs to be re-run on a 'single' processor.
	// Check that all the command-line options are set.
	if (missedRank >= 0 || missedProcs >= 0) {
		BOOL exitEarly = FALSE;

		if (numprocs != 1) exitEarly = TRUE;
		else if (missedRank == -1 || missedProcs == -1) exitEarly = TRUE;
		else if (!(missedRank >= 0 && missedRank <= missedProcs)) exitEarly = TRUE;

		if (exitEarly) {
			MPI_Finalize();
			return -1;
		} else {
			myrank = missedRank;
			numprocs = missedProcs;
		}
	} 

	if (!isDoingMerge) {
		// Do the non-permuted association calculation across all ranks 
		// in the cluster allocated to this job.
		if ((map = bmapfile_load(&ops, &nMap)) != NULL && 
			(samples = samplefile_load(&ops, &nSamples)) != NULL &&
			(pPermutations = permfile_load(&ops, &perm_rows, &perm_columns)) != NULL &&
			(spedData = spedfile_load(&ops, &nSNPs, &nSpedSamples, &nSpedRecordSizeInBytes)) != NULL)
		{
			if (nMap > 0 && perm_rows > 0 && perm_columns == nSamples && nSNPs == nMap) {
				int nPermsPerProcessor = get_tasks_per_processor(numprocs, perm_rows);
				int start_perm, end_perm, i;
				BOOL adjAllocated = TRUE;

				start_perm = myrank * nPermsPerProcessor;
				end_perm = start_perm + nPermsPerProcessor;
				if (end_perm > perm_rows || myrank == numprocs - 1) end_perm = perm_rows;

				originals = create_vector(nSNPs, 0);
				results = create_vector(nSNPs, 0);
				R = create_int_array(nSNPs, 0);
				maxR = create_int_array(nSNPs, 0);
				if (myrank == ROOT) {
					adj = adjusted_create_array(nSNPs);
					if (adj == NULL) adjAllocated = FALSE;
				}

				// Permute the data-set.
				if (originals != NULL && R != NULL && maxR != NULL && results != NULL) {
					samplefile_aff_coding(samples, nSamples);
					samplefile_set_original_order(samples, nSamples);

					// Non-permuted calculation.
					for (i = 0; i < nSNPs; i++) {
						BOOL X = FALSE, haploid = FALSE;
						BitString bs = (BitString) (spedData + (i * nSpedRecordSizeInBytes));

						// Check if the SNP is located on autosome.
						if (chromos.chr_sex[map[i].chr]) X = TRUE;
						else if (chromos.chr_haploid[map[i].chr]) haploid = TRUE;

						// Calculate the original pvalue and affected counts.
						originals[i] = pmodel_assoc(X, haploid, bs, samples, nSamples, &ops, NULL, NULL, 0, 0, NULL, 0);
						if (myrank == ROOT) adj[i].chisq = originals[i];
					}

					for (i = start_perm; i < end_perm && i < perm_rows; i++) {
						long offset = (long) i * (long) perm_columns;
						int *pPerms = pPermutations + offset;

						zero_vector(results, nSNPs);
						samplefile_permute_phenotype(samples, pPerms, nSamples);
						for (j = 0; j < nSNPs; j++) {
							BOOL X = FALSE, haploid = FALSE;
							BitString bs = (BitString) (spedData + (j * nSpedRecordSizeInBytes));

							// Check if the SNP is located on autosome.
							if (chromos.chr_sex[map[j].chr]) X = TRUE;
							else if (chromos.chr_haploid[map[j].chr]) haploid = TRUE;

							// Calculate the original pvalue and affected counts.
							results[j] = pmodel_assoc(X, haploid, bs, samples, nSamples, &ops, NULL, NULL, 0, 0, NULL, 0);
						}
						perm_update(nSNPs, results, originals, R, maxR);
					}

					// Dump the local counts to file.
					jobDone = count_file_write_and_copy(&ops, SUFFIX_MODEL_FILE, myrank, numprocs, abs(end_perm - start_perm), nSNPs, R, maxR) ? TRUE : FALSE;
					if (myrank == ROOT && jobDone) {
						jobDone = FALSE;
						if (pmodel_write_model_options(&ops)) jobDone = adjusted_write_and_copy(adj, nSNPs, &ops, myrank) ? TRUE : FALSE;
					}
				}
			}
		}
	} else { 
		// Merging & sorting the output files.
		if (numprocs == 1) { // This code in this mode needs to run as a singleton.
			char line[MAX_LINE_LENGTH];
			BOOL *file_copy_failures = NULL, hasCopiedAll = TRUE, hasReadModelOptions = FALSE;

			map = _bmapfile_load(&ops, &nMap);
			spedData = _spedfile_load(&ops, &nSNPs, &nSpedSamples, &nSpedRecordSizeInBytes);
			samples = _samplefile_load(&ops, &nSamples);
			numprocs = readNumberOfProcs(remoteFilepath, line, MAX_LINE_LENGTH, ops.rdir, ops.szRunId, SUFFIX_MODEL_FILE);
			hasReadModelOptions = pmodel_read_model_options(&ops) ? TRUE : FALSE;
			adj = _adjusted_load(&ops, &nAdj);

			if (map != NULL && nMap > 0 && spedData != NULL && nSNPs > 0 && nSNPs == nMap && samples != NULL && nSamples > 0 && numprocs > 0 &&
				adj != NULL && nAdj > 0 && nAdj == nSNPs) 
			{
				int nPermutations = 0;
				file_copy_failures = create_bool_array(numprocs, FALSE); 
				R = create_int_array(nSNPs, 0);
				maxR = create_int_array(nSNPs, 0);
				if (count_file_copy_and_parse(&ops, SUFFIX_MODEL_FILE, numprocs, nSNPs, &nPermutations, R, maxR, &hasCopiedAll, file_copy_failures)) {
					if (hasCopiedAll) {
						FILE *out = NULL;

						sprintf(localFilepath, "%s%s%s", ops.ldir, ops.szRunId, SUFFIX_MODEL_FILE);
						sprintf(remoteFilepath, "%s%s%s", ops.rdir, ops.szRunId, SUFFIX_MODEL_FILE);
						out = fopen(localFilepath, "w");
						if (out != NULL) {
							double nperms = (double) nPermutations + 1;

							samplefile_aff_coding(samples, nSamples);
							samplefile_set_original_order(samples, nSamples);

							// Write out the results of each SNP to file.
							pmodel_write_file_header(out, &ops);
							for (i = 0; i < nSNPs; i++) {
								double emp1, emp2;
								BOOL X = FALSE, haploid = FALSE;
								BitString bs = (BitString) (spedData + (i * nSpedRecordSizeInBytes));

								// Check if the SNP is located on autosome.
								if (chromos.chr_sex[map[i].chr]) X = TRUE;
								else if (chromos.chr_haploid[map[i].chr]) haploid = TRUE;

								// Ratio counts
								emp1 = (double) (R[i] + 1) / nperms;
								emp2 = (double) (maxR[i] + 1) / nperms;

								pmodel_assoc(X, haploid, bs, samples, nSamples, &ops, &map[i], out, emp1, emp2, &adj[i], zt);
							}
							fclose(out);
						}

						// Final sort and clean up.
						jobDone = sort_file_by_col_and_copy(localFilepath, remoteFilepath, ops.szSortColumnName) ? TRUE : FALSE;	
						if (jobDone) assoc_cleanup_temp_files(&ops, SUFFIX_MODEL_FILE, numprocs);
					}
				} else {	
					write_failed_rank_log(&ops, numprocs, file_copy_failures);
				}
				if (file_copy_failures != NULL) free(file_copy_failures);
			}
		}
	}

	// Free now as not needed anymore
	if (pPermutations != NULL) free(pPermutations);
	if (samples != NULL) free(samples);
	if (spedData != NULL) free(spedData);
	if (map != NULL) free(map); 
	if (originals != NULL) free(originals);
	if (R != NULL) free(R);
	if (maxR != NULL) free(maxR);
	if (results != NULL) free(results);
	if (adj != NULL) free(adj);
	chromosomes_free_attributes(&chromos); 

	MPI_Finalize();

	if (jobDone) return 0;
	else return -1;
}

int main(int argc, char *argv[]) {
	return pmodel_main(argc, argv);
}